#!/usr/bin/env python
# -*- coding: utf-8 -*-
from flask import Blueprint, request, current_app, g
from applications.extensions.db_plugin import db
from applications.extensions.jwt_plugin import jwt_required_with_user
from common.utils.terminal_exec import exec_shell
from applications.extensions.redis_plugin import redis_client
from common.core.make_response import (
    make_success, make_error, make_validation_error, make_forbidden
)
from common.settings.constants import (
    DATE_FORMAT_COMPACT, DATETIME_FORMAT_SHORT_DASH,
    DATE_FORMAT_DASH, DATETIME_FORMAT_SHORT_COMPACT
)
from applications.models.terminal import FlowmeterDailyStats
from datetime import datetime, timedelta
import re, json

terminal_bp = Blueprint('terminal', __name__, url_prefix='/api/terminal')


def is_valid_port(port):
    """检查端口是否合法"""
    return isinstance(port, int) and 1 <= port <= 65535


def iptables_rule_exists(port):
    """检查该端口是否已有规则"""
    result = exec_shell("iptables -L OUTPUT -n --line-numbers")
    if not result["success"]:
        return False  # 查询失败时默认继续
    return f"spt:{port}" in result["stdout"]


@terminal_bp.route('/iptables/get', methods=['GET'])
def iptables_get():
    """
    获取当前 iptables OUTPUT 规则，将解析的结构化数据放到Redis
    :return:
    """
    secret_key = request.args.get('secret_key', '').strip()
    current_app.logger.info(f"secret_key: {request.args}")

    if secret_key != current_app.config.get("SECRET_KEY"):
        return make_forbidden(msg="SECRET_KEY无效")

    result = exec_shell("iptables -L OUTPUT -v -n -x --line-number")

    if not result["success"]:
        return make_error(msg=result["stderr"])

    # 解析成结构化数据
    rules = []

    for line in result['stdout'].split("\n"):
        j = [field for field in line.split(" ") if field]
        if not j or not j[0].isnumeric() or len(j) < 3:
            continue  # 跳过无效行
        try:
            pkts_used = int(j[1])
            bytes_used = int(j[2])
        except ValueError:
            continue  # 流量字段无效时跳过
        gb = round(bytes_used / (1024 ** 3), 4)  # 优化单位计算
        # 使用正则表达式从字符串末尾匹配端口号：1. 检查字段末尾是否含":数字"格式；2. 提取端口号保存至port变量；3. 匹配失败则跳过该行
        port_match = re.search(r':(\d+)$', j[-1])
        if not port_match:
            continue  # 端口格式异常时跳过
        port = port_match.group(1)

        rule = {
            "port": port,
            "gb": gb,
            "pkts": pkts_used,
            "bytes": bytes_used,
        }

        rules.append(rule)

    y_m_d_h_m = datetime.now().strftime(DATETIME_FORMAT_SHORT_COMPACT)  # 当前时间 年月日时分
    y_m_d = datetime.now().strftime(DATE_FORMAT_COMPACT)  # 当前时间 年月日
    # 保存到 Redis，使用哈希表存储，key 为y_m_d_h_m，value 为规则，默认使用 JSON 格式保存，有效期90天
    try:
        redis_client.set_hash(
            name=f"iptables:output:{y_m_d}",
            key=y_m_d_h_m,
            value=rules,
            is_json=True,
            time=2592000
        )
    except Exception as e:
        current_app.logger.error(f"Failed to write to Redis: {e}")
        return make_error(msg="Failed to save data to Redis")

    return make_success(data=rules)


def process_iptables_ports(secret_key, ports, action='set', date=None):
    """
    添加或清除端口规则
    action: 'set' -> 添加, 'clear' -> 删除, 'save' -> 保存流量数据到数据库
    """
    expected_secret = current_app.config.get("SECRET_KEY")

    # 验证密钥
    if not isinstance(secret_key, str) or secret_key != expected_secret:
        return make_forbidden(msg="SECRET_KEY无效")

    if not isinstance(ports, list):
        return make_validation_error(msg="请传入端口号列表（ports: List[int]）")

    for port in ports:
        if not isinstance(port, int):
            return make_validation_error(msg=f"端口必须为整数，非法端口: {port}")

    if action not in ('set', 'clear', 'save'):
        return make_validation_error(msg="action 参数必须为 'set'、'clear'、'save'")

    # 记录执行结果
    summary = {"success": [], "failed": [], "skipped": []}

    if action == 'save':
        target_date = datetime.now().strftime(DATE_FORMAT_COMPACT)
        if date:
            target_date = date
        try:
            for port in ports:
                if not is_valid_port(port):
                    summary["failed"].append({"port": port, "error": "非法端口"})
                    continue
                try:
                    result, total_gb, total_pkts, total_bytes = get_hourly_data_for_date(str(port), target_date)
                    stat_date = datetime.strptime(target_date, DATE_FORMAT_COMPACT).date()
                    existing_record = FlowmeterDailyStats.query.filter_by(
                        port=port,
                        stat_date=stat_date
                    ).first()

                    if existing_record:
                        existing_record.total_gb = total_gb
                        existing_record.total_pkts = total_pkts
                        existing_record.total_bytes = total_bytes
                        db.session.add(existing_record)
                    else:
                        new_record = FlowmeterDailyStats(
                            port=port,
                            stat_date=stat_date,
                            total_gb=total_gb,
                            total_pkts=total_pkts,
                            total_bytes=total_bytes
                        )
                        db.session.add(new_record)

                    summary["success"].append({
                        "port": port,
                        "total_gb": total_gb,
                        "total_pkts": total_pkts,
                        "total_bytes": total_bytes
                    })
                except Exception as e:
                    summary["failed"].append({"port": port, "error": str(e)})

            db.session.commit()
        except Exception as e:
            db.session.rollback()
            return make_validation_error(msg=f"数据库操作失败：{e}")
        return make_success(data=summary)

    # set 和 clear 操作
    for port in ports:
        if not is_valid_port(port):
            summary["failed"].append({"port": port, "error": "非法端口"})
            continue

        exists = iptables_rule_exists(port)

        if action == 'set':
            if exists:
                summary["skipped"].append({"port": port, "reason": "规则已存在"})
                continue
            cmd = ["iptables", "-A", "OUTPUT", "-p", "tcp", "--sport", str(port)]
        else:  # clear
            if not exists:
                summary["skipped"].append({"port": port, "reason": "规则不存在"})
                continue
            cmd = ["iptables", "-D", "OUTPUT", "-p", "tcp", "--sport", str(port)]

        try:
            result = exec_shell(cmd)
        except Exception as e:
            summary["failed"].append({"port": port, "error": str(e)})
            continue

        if result["success"]:
            summary["success"].append(port)
        else:
            summary["failed"].append({"port": port, "error": result["stderr"]})

    return make_success(data=summary)


@terminal_bp.route('/iptables/set', methods=['POST'])
def iptables_set():
    """
    添加端口规则
    """
    secret_key = request.args.get('secret_key', '').strip()
    data = request.get_json()
    return process_iptables_ports(
        secret_key=secret_key,
        ports=data.get("ports", []),
        action="set"
    )


@terminal_bp.route('/iptables/clear', methods=['POST'])
def iptables_clear():
    """
    删除端口规则
    """
    secret_key = request.args.get('secret_key', '').strip()
    data = request.get_json()
    return process_iptables_ports(
        secret_key=secret_key,
        ports=data.get("ports", []),
        action="clear"
    )


@terminal_bp.route('/iptables/save', methods=['POST'])
def iptables_save():
    """
    保存 iptables 规则，保存当天所有的端口流量到数据库
    """
    secret_key = request.args.get('secret_key', '').strip()
    data = request.get_json()
    return process_iptables_ports(
        secret_key=secret_key,
        ports=data.get("ports", []),
        date=data.get("date", None),
        action="save"
    )


@terminal_bp.route('/iptables/save_db', methods=['POST'])
def iptables_save_db():
    """
    手动保存流量数据，支持单个或批量
    {...}、[{...}, {...}]
    :return:
    """
    try:
        secret_key = request.args.get('secret_key', '').strip()
        data_list = request.get_json()

        if secret_key != current_app.config.get("SECRET_KEY"):
            return make_forbidden(msg="SECRET_KEY无效")

        # 如果传入的是单个对象，则转为列表
        if isinstance(data_list, dict):
            data_list = [data_list]
        elif not isinstance(data_list, list):
            return make_validation_error(msg="请求体必须是数组或单个对象")

        records_to_insert = []
        records_to_update = []
        response_data = {"success": [], "failed": []}

        for idx, data in enumerate(data_list):
            try:
                port = data.get("port")
                target_date = data.get("stat_date")
                total_bytes = data.get("total_bytes")
                total_pkts = data.get("total_pkts")
                total_gb = data.get("total_gb")

                # 验证必填字段
                if None in [port, target_date, total_bytes, total_pkts, total_gb]:
                    response_data["failed"].append({
                        "index": idx,
                        "error": "缺少必要字段"
                    })
                    continue

                stat_date = datetime.strptime(target_date, DATE_FORMAT_COMPACT).date()

                # 查询是否已存在
                existing_record = FlowmeterDailyStats.query.filter_by(
                    port=port,
                    stat_date=stat_date
                ).first()

                if existing_record:
                    # 更新现有记录
                    existing_record.total_gb = total_gb
                    existing_record.total_pkts = total_pkts
                    existing_record.total_bytes = total_bytes
                    records_to_update.append(existing_record)
                    response_data["success"].append({
                        "index": idx,
                        "status": "updated",
                        "port": port,
                        "stat_date": target_date
                    })
                else:
                    # 准备插入新记录
                    new_record = FlowmeterDailyStats(
                        port=port,
                        stat_date=stat_date,
                        total_gb=total_gb,
                        total_pkts=total_pkts,
                        total_bytes=total_bytes
                    )
                    records_to_insert.append(new_record)
                    response_data["success"].append({
                        "index": idx,
                        "status": "inserted",
                        "port": port,
                        "stat_date": target_date
                    })

            except ValueError as ve:
                response_data["failed"].append({
                    "index": idx,
                    "error": f"日期格式错误: {ve}"
                })
            except Exception as e:
                response_data["failed"].append({
                    "index": idx,
                    "error": str(e)
                })

        # 批量插入和更新
        try:
            if records_to_insert:
                db.session.add_all(records_to_insert)
            if records_to_update:
                # SQLAlchemy 会自动跟踪 dirty state，无需额外操作
                pass
            db.session.commit()
            return make_success(data=response_data)
        except Exception as e:
            db.session.rollback()
            return make_error(msg=f"数据库操作失败: {e}")

    except Exception as e:
        return make_error(msg=f"请求处理失败: {e}")


def get_database_summary_data(port, days):
    """
    从数据库获取指定端口最近N天的每日汇总数据
    :param port: 端口号
    :param days: 天数
    :return: (result_list, total_gb)
    """
    result = []
    total_gb = 0.0

    end_date = datetime.now().date()
    start_date = end_date - timedelta(days=days - 1)  # 包括今天在内共days天

    # 从数据库查询最近N天的数据
    records = FlowmeterDailyStats.query.filter(
        FlowmeterDailyStats.port == int(port),
        FlowmeterDailyStats.stat_date.between(start_date, end_date)
    ).order_by(FlowmeterDailyStats.stat_date.asc()).all()

    for record in records:
        result.append({
            "time": record.stat_date.strftime(DATE_FORMAT_DASH),
            "gb": float(record.total_gb),
            "pkts": int(record.total_pkts),
            "bytes": int(record.total_bytes)
        })
        total_gb += float(record.total_gb)

    return result, total_gb


def get_hourly_data_for_date(port, target_date):
    """
    获取指定端口在指定日期的每小时流量数据（30天内的）
    :param port: 端口号
    :param target_date: 日期，格式为 YYYYMMDD
    :return: (result_list, total_gb, total_pkts, total_bytes)
    """
    result = []
    total_gb = 0.0
    total_pkts = 0
    total_bytes = 0

    redis_key = f"iptables:output:{target_date}"
    fields = redis_client.get_client().hkeys(redis_key)

    current_app.logger.info(f"查询redis数据：{redis_key}")

    for field in fields:
        try:
            field_str = field.decode() if isinstance(field, bytes) else field
            ts = datetime.strptime(field_str, DATETIME_FORMAT_SHORT_COMPACT)

            raw_json = redis_client.get_hash(redis_key, field_str)

            if not raw_json:
                continue

            if isinstance(raw_json, bytes):
                raw_json = raw_json.decode('utf-8')

            data_list = json.loads(raw_json)
            for item in data_list:
                if str(item.get('port')) == port:
                    gb_val = item.get('gb', 0)
                    pkts_val = item.get('pkts', 0)
                    bytes_val = item.get('bytes', 0)

                    # 累加所有三个指标
                    total_gb += float(gb_val)
                    total_pkts += int(pkts_val)
                    total_bytes += int(bytes_val)

                    result.append({
                        "time": ts.strftime(DATETIME_FORMAT_SHORT_DASH),
                        "gb": gb_val,
                        "pkts": pkts_val,
                        "bytes": bytes_val
                    })
                    break
        except Exception as sub_e:
            current_app.logger.warning(f"解析字段失败: {field} - {sub_e}")
            continue

    current_app.logger.info(f"查询结果：total_gb[{total_gb}],total_pkts[{total_pkts}],total_bytes[{total_bytes}]")
    return result, total_gb, total_pkts, total_bytes


def get_database_range_data(port, start_date, end_date):
    """
    从数据库获取指定端口在指定日期范围内的每日汇总数据
    :param port: 端口号
    :param start_date: 起始日期
    :param end_date: 结束日期
    :return: (result_list, total_gb)
    """
    result = []
    total_gb = 0.0

    # 从数据库查询日期范围内的数据
    records = FlowmeterDailyStats.query.filter(
        FlowmeterDailyStats.port == int(port),
        FlowmeterDailyStats.stat_date.between(start_date, end_date)
    ).order_by(FlowmeterDailyStats.stat_date.asc()).all()

    for record in records:
        result.append({
            "time": record.stat_date.strftime(DATE_FORMAT_DASH),
            "gb": float(record.total_gb),
            "pkts": int(record.total_pkts),
            "bytes": int(record.total_bytes)
        })
        total_gb += float(record.total_gb)

    return result, total_gb


@terminal_bp.route('/flowmeter/search', methods=['GET'])
@jwt_required_with_user()
def flowmeter_search():
    """
    查询指定端口在当天的流量数据（Redis）
    默认查询当天每小时数据，支持查询最近7天或30天的每日汇总数据，也可查询指定日期数据
    新增支持按时间范围查询（start_date - end_date）
    """
    port = request.args.get('port', '').strip()
    query_type = request.args.get('type', 'day').strip()  # day (默认) or week or month or range
    date = request.args.get('date', '').strip()  # 查询指定日期，格式: YYYYMMDD
    start_date = request.args.get('start_date', '').strip()  # 起始日期，格式: YYYYMMDD
    end_date = request.args.get('end_date', '').strip()  # 结束日期，格式: YYYYMMDD

    current_app.logger.info(
        f"[查询端口流量数据] port: {port}, query_type: {query_type}, "
        f"date: {date}, start_date: {start_date}, end_date: {end_date}"
    )

    if not port:
        return make_validation_error(msg='缺少端口参数')

    try:
        if query_type == 'range':
            # 按时间范围查询
            if not start_date or not end_date:
                return make_validation_error(msg='时间范围查询需要提供 start_date 和 end_date 参数')

            try:
                start_dt = datetime.strptime(start_date, DATE_FORMAT_COMPACT).date()
                end_dt = datetime.strptime(end_date, DATE_FORMAT_COMPACT).date()

                if start_dt > end_dt:
                    return make_validation_error(msg='起始日期不能晚于结束日期')

                # 查询范围内的每日汇总数据
                result, total_gb = get_database_range_data(port, start_dt, end_dt)
            except ValueError:
                return make_validation_error(msg='日期格式错误，应为 YYYYMMDD')

        elif query_type == 'day':
            # 默认查询当天每小时数据，如果指定日期则查询指定日期数据
            if not date:
                target_date = datetime.now().strftime(DATE_FORMAT_COMPACT)
            else:
                try:
                    datetime.strptime(date, DATE_FORMAT_COMPACT)  # 验证日期格式
                    target_date = date
                except ValueError:
                    return make_validation_error(msg='日期格式错误，应为 YYYYMMDD')

            result, total_gb, total_pkts, total_bytes = get_hourly_data_for_date(port, target_date)

        elif query_type == 'week':  # 查询最近7天数据（按天汇总）
            result, total_gb = get_database_summary_data(port, days=7)

        elif query_type == 'month':  # 查询最近30天数据（按天汇总）
            result, total_gb = get_database_summary_data(port, days=30)

        else:
            return make_validation_error(msg='type 参数必须为 day、week、month 或 range')

    except Exception as e:
        current_app.logger.exception(f"查询失败: {e}")
        return make_error(msg='查询失败，请稍后重试')

    result.sort(key=lambda x: x['time'])

    # 返回数据加个 total_gb 字段，不影响 data 结构
    return make_success(data=result, extra={"total_gb": round(total_gb, 4)})
